import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve
import os


def MFEABernoulli3Ddata(filename):
    Node = []
    Beam = []
    Spring = []
    MaterialBeam = []
    MaterialSpring = []
    Load = []
    Constr = []
    DF = []

    isNodeSection = False
    isBeamSection = False
    isSpringSection = False
    isMaterialBeamSection = False
    isMaterialSpringSection = False
    isLoadSection = False
    isDFSection = False
    isConstrSection = False

    with open(os.path.join("input", f"{filename}.txt"), 'r') as fid:
        for line in fid:
            line = line.strip()
            if not line:
                continue

            if line.startswith('*Node'):
                isNodeSection = True
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*Beam'):
                isNodeSection = False
                isBeamSection = True
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*Spring'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = True
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*MaterialBeam'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = True
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*MaterialSpring'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = True
                isLoadSection = False
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*Load'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = True
                isDFSection = False
                isConstrSection = False
                continue
            elif line.startswith('*DF'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = True
                isConstrSection = False
                continue
            elif line.startswith('*Constr'):
                isNodeSection = False
                isBeamSection = False
                isSpringSection = False
                isMaterialBeamSection = False
                isMaterialSpringSection = False
                isLoadSection = False
                isDFSection = False
                isConstrSection = True
                continue

            if isNodeSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 4:
                    Node.append(data[1:])
            elif isBeamSection:
                data = list(map(int, line.split(',')))
                if len(data) >= 3:
                    Beam.append(data[1:])
            elif isSpringSection:
                data = list(map(int, line.split(',')))
                if len(data) >= 3:
                    Spring.append(data[1:])
            elif isMaterialBeamSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 6:
                    MaterialBeam.append(data[:6])
            elif isMaterialSpringSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 6:
                    MaterialSpring.append(data[:6])
            elif isLoadSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 3:
                    Load.append([int(data[0]), data[1], data[2]])
            elif isDFSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 1:
                    DF.append(data)
            elif isConstrSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 3:
                    Constr.append([int(data[0]), int(data[1]), data[2]])
    return np.array(Node), np.array(Beam), np.array(Spring), np.array(MaterialBeam), np.array(MaterialSpring), np.array(
        Load), np.array(Constr)

def equivalent_node_force(node, element):
    # 计算线性分布荷载的等效节点力
    a = Load[:, 0]
    b = Load[:, 1]
    c = Load[:, 2]
    enf = np.zeros(12)  # 定义12x1的等效节点力向量
    x = node[:, 0]
    y = node[:, 1]
    z = node[:, 2]
    bar_length = BarsLength(x, y, z,element)
    EleCountBeam = element.shape[0]
    for iEle in range(EleCountBeam):
        p1 = b[iEle - 1]
        p2 = c[iEle - 1]
        idof = a[iEle - 1]
        L = bar_length[iEle]
        n1 = element[iEle, 0]-1
        n2 = element[iEle, 1]-1
        R = CoordTransform([x[n1], x[n2]], [y[n1], y[n2]],[z[n1], z[n2]], L)

        if idof == 1:  # 分布轴向力
            enf[0] = (2 * p1 + p2) * L / 6
            enf[7] = (p1 + 2 * p2) * L / 6
        elif idof == 2:  # 分布横向力
            enf[1] = (7 * p1 + 3 * p2) * L / 20
            enf[3] = (3 * p1 + 2 * p2) * L**2 / 60
            enf[7] = (3 * p1 + 7 * p2) * L / 20
            enf[9] = -(2 * p1 + 3 * p2) * L**2 / 60
        elif idof == 3:  # 分布弯矩
            enf[1] = -(p1 + p2) / 2
            enf[2] = (p1 - p2) * L / 12
            enf[4] = (p1 + p2) / 2
            enf[5] = -(p1 - p2) * L / 12
        elif idof == 4:  # 集中力  p1为距左节点距离，p2为集中力大小
            p3 = bar_length[iEle]-p1
            enf[1] = -(p2*(p3**2) /(bar_length[iEle])**3) *(3*p1 + p3)
            enf[7] = -(p2*(p1**2) /(bar_length[iEle])**3) *(p1 + 3*p3)
            enf[4] = -p2*p1*(p3**2)/(bar_length[iEle]**2)
            enf[10] = p2*p3*(p1**2)/(bar_length[iEle]**2)

        enf = np.dot(R.T, enf)  # 把等效节点力转换到整体坐标下
        Dof = 6
        Dofs = Dof * Node.shape[0]
        FFG = np.zeros((Dofs, 1))
        n1_dofs = np.arange(n1 * Dof, (n1 + 1) * Dof)
        n2_dofs = np.arange(n2 * Dof, (n2 + 1) * Dof)
        ElementNodeDOFBeam = np.concatenate((n1_dofs, n2_dofs))
        for i in range(len(ElementNodeDOFBeam)):
            FFG[ElementNodeDOFBeam[i]] += enf[i]
    print(FFG)
    return FFG



def BarsLength(x, y, z, ele):
    return np.sqrt((x[ele[:, 1] - 1] - x[ele[:, 0] - 1]) ** 2 +
                   (y[ele[:, 1] - 1] - y[ele[:, 0] - 1]) ** 2 +
                   (z[ele[:, 1] - 1] - z[ele[:, 0] - 1]) ** 2)


def CoordTransform(x, y, z, L):
    x1, x2 = x
    y1, y2 = y
    z1, z2 = z
    if x1 == x2 and y1 == y2:
        if z2 > z1:
            Lambda = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
        else:
            Lambda = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]])
    else:
        CXx = (x2 - x1) / L
        CYx = (y2 - y1) / L
        CZx = (z2 - z1) / L
        D = np.sqrt(CXx * CXx + CYx * CYx)
        CXy = -CYx / D
        CYy = CXx / D
        CZy = 0
        CXz = -CXx * CZx / D
        CYz = -CYx * CZx / D
        CZz = D
        Lambda = np.array([[CXx, CYx, CZx],
                           [CXy, CYy, CZy],
                           [CXz, CYz, CZz]])
    R = np.block([
        [Lambda, np.zeros((3, 9))],
        [np.zeros((3, 3)), Lambda, np.zeros((3, 6))],
        [np.zeros((3, 6)), Lambda, np.zeros((3, 3))],
        [np.zeros((3, 9)), Lambda]
    ])
    return R


if __name__ == '__main__':
    filename = 'zaihe'
    Node, Beam, Spring, MaterialBeam, MaterialSpring, Load, Constr = MFEABernoulli3Ddata(filename)
    combined_ENF=equivalent_node_force(Node, Beam)